unstract/sdk1/tests/patches/test_litellm_cohere_timeout.py (1)

7-10: Exercise the production import hook in one test.

Because this file imports the patch module directly, the wiring checks only prove direct-import behavior. Add one integration test that imports unstract.sdk1.embedding and asserts the handler is patched, so Line 7 in unstract/sdk1/src/unstract/sdk1/embedding.py cannot regress silently.

Also applies to: 194-203

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/tests/patches/test_litellm_cohere_timeout.py` around lines 7 -
10, Add an integration test that imports unstract.sdk1.embedding (instead of
importing the patch module directly) and asserts the production import hook
applied the patch: after importing unstract.sdk1.embedding, verify the module's
embedding handler uses the patched implementations by checking that its async
and sync embedding callables resolve to the symbols _patched_async_embedding and
_patched_embedding (or that their identities/reference equality match those
functions imported from unstract.sdk1.patches.litellm_cohere_timeout); update or
add this assertion into test_litellm_cohere_timeout.py alongside the existing
direct-import checks so the production import wiring (line referencing
unstract.sdk1.embedding) is exercised and cannot regress.

unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py (1)

26-64: ⚠️ Potential issue | 🟠 Major

Move the private LiteLLM imports behind the guard and gate on the exact tested version.

Lines 26-43 still import private litellm.llms.* modules before _SKIP_PATCH is evaluated, so a future LiteLLM reorg can fail at import time instead of cleanly skipping. Also, Line 56 only skips > 1.80.0; for a copied private implementation, != is safer. The skip warning on Lines 58-64 uses DeprecationWarning, which is usually hidden for imported library code, so operators may never notice that the patch was not applied. With the provided dependency context showing LiteLLM 1.81.7, this path would currently skip.

🔧 Suggested fix

 import importlib.metadata
 import json
 import logging
 import warnings
 from collections.abc import Callable

 import httpx
 import litellm
-import litellm.llms.bedrock.embed.embedding as _bedrock_embed
-import litellm.llms.cohere.embed.handler as _cohere_handler
 from litellm.litellm_core_utils.litellm_logging import (
     Logging as LiteLLMLoggingObj,
 )
-from litellm.llms.cohere.embed.handler import (
-    validate_environment,
-)
-from litellm.llms.cohere.embed.v1_transformation import (
-    CohereEmbeddingConfig,
-)
-from litellm.llms.custom_httpx.http_handler import (
-    AsyncHTTPHandler,
-    HTTPHandler,
-    get_async_httpx_client,
-)
-from litellm.types.llms.bedrock import CohereEmbeddingRequest
-from litellm.types.utils import EmbeddingResponse
 from packaging.version import Version

 logger = logging.getLogger(__name__)

 _DEFAULT_TIMEOUT = httpx.Timeout(None)

 _PATCHED_LITELLM_VERSION = "1.80.0"
 _litellm_version = importlib.metadata.version("litellm")
-_SKIP_PATCH = Version(_litellm_version) > Version(_PATCHED_LITELLM_VERSION)
+_SKIP_PATCH = Version(_litellm_version) != Version(_PATCHED_LITELLM_VERSION)
 if _SKIP_PATCH:
     warnings.warn(
         "litellm_cohere_timeout patch was SKIPPED — not applied. "
         f"Current litellm version: {_litellm_version}. "
         f"Patch was written for: {_PATCHED_LITELLM_VERSION}. "
         "Please verify the upstream fix and remove this module.",
-        DeprecationWarning,
+        RuntimeWarning,
         stacklevel=2,
     )
+else:
+    import litellm.llms.bedrock.embed.embedding as _bedrock_embed
+    import litellm.llms.cohere.embed.handler as _cohere_handler
+    from litellm.llms.cohere.embed.handler import validate_environment
+    from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig
+    from litellm.llms.custom_httpx.http_handler import (
+        AsyncHTTPHandler,
+        HTTPHandler,
+        get_async_httpx_client,
+    )
+    from litellm.types.llms.bedrock import CohereEmbeddingRequest
+    from litellm.types.utils import EmbeddingResponse

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py` around
lines 26 - 64, The current file imports private litellm modules before computing
_SKIP_PATCH, which can cause import-time failures; move all imports of
litellm.llms.* and litellm.litellm_core_utils.* (symbols: _bedrock_embed,
_cohere_handler, LiteLLMLoggingObj, validate_environment, CohereEmbeddingConfig,
AsyncHTTPHandler, HTTPHandler, get_async_httpx_client, CohereEmbeddingRequest,
EmbeddingResponse) so they occur only after computing _litellm_version and
_SKIP_PATCH; change the version check to use equality against
_PATCHED_LITELLM_VERSION (i.e. skip only when Version(_litellm_version) ==
Version(_PATCHED_LITELLM_VERSION) or gate with != as requested) and replace the
DeprecationWarning with a visible warning class such as UserWarning or
RuntimeWarning so operators see the warning when the patch is skipped.

unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py (1)

15-15: Consider adding a TODO marker for discoverability.

Adding TODO: makes this line searchable via grep/IDE tooling, helping ensure timely removal once upstream fixes the bug.

-Remove this patch when litellm ships a fix upstream.
+TODO: Remove this patch when litellm ships a fix upstream.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py` at line
15, Update the top-of-file note in litellm_cohere_timeout.py so it includes a
searchable TODO marker; specifically replace or edit the comment "Remove this
patch when litellm ships a fix upstream." to start with "TODO:" (e.g., "TODO:
Remove this patch when litellm ships a fix upstream.") so it is discoverable via
grep/IDE and clearly signals removal once the upstream bug in litellm is fixed.

unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py (1)

48-65: ⚠️ Potential issue | 🔴 Critical

Import the modules you rebind later.

This block imports symbols from the Cohere and Bedrock modules, but lines 213-215 assign through _cohere_handler and _bedrock_embed, which are never defined. On LiteLLM 1.81.7, importing this patch will raise NameError before the monkey-patch is applied, breaking the side-effect activation path from unstract.sdk1.embedding.

🐛 Proposed fix

     import httpx
     import litellm
+    import litellm.llms.bedrock.embed.embedding as _bedrock_embed
+    import litellm.llms.cohere.embed.handler as _cohere_handler
     from litellm.litellm_core_utils.litellm_logging import (
         Logging as LiteLLMLoggingObj,
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@unstract/sdk1/src/unstract/sdk1/patches/litellm_cohere_timeout.py` around
lines 48 - 65, The patch imports Cohere/Bedrock symbols but later rebinds
globals like _cohere_handler and _bedrock_embed without first defining or
importing them, causing a NameError on import; to fix, import or define the
original symbols before rebinding (e.g., import the existing cohere handler and
bedrock embed functions used in lines that assign to _cohere_handler and
_bedrock_embed), then perform the monkey-patch assignments; locate the rebinding
logic targeting _cohere_handler and _bedrock_embed and ensure the original
symbols are referenced (via imports or safe getattr fallbacks) so the module can
be imported without raising NameError.